import xgboost as xgb
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.metrics import root_mean_squared_error
import pandas as pd
import matplotlib.pyplot as plt

"""
XGBoost回归示例 - 使用糖尿病数据集
糖尿病数据集包含10个特征，目标变量是糖尿病进展的定量测量
特征数：10个（年龄、性别、BMI、血压等）
样本数：442个
目标变量：糖尿病进展的定量测量（连续值）
"""

# 使用糖尿病数据集（无需网络下载）
# load_diabetes() 函数加载内置的糖尿病数据集
# return_X_y=True 参数表示直接返回特征矩阵X和目标变量y
X, y = load_diabetes(return_X_y=True)  # y 是糖尿病进展的定量测量

# 展示数据信息
print("X的数据类型:", type(X))  # numpy.ndarray
print("y的数据类型:", type(y))  # numpy.ndarray
print("X的形状:", X.shape)
print("y的形状:", y.shape)
print("目标变量统计:")
print(f"  最小值: {y.min():.2f}")
print(f"  最大值: {y.max():.2f}")
print(f"  均值: {y.mean():.2f}")
print(f"  标准差: {y.std():.2f}")

# 划分数据集
# train_test_split() 函数将数据集划分为训练集和测试集
# test_size=0.2 表示测试集占总数据的20%，训练集占80%
# random_state=42 设置随机种子，确保每次划分结果一致，便于结果复现
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)


def train_model(X_train, y_train, max_depth=3, n_estimators=100, show_progress=True):
    """
    训练XGBoost回归模型
    :param X_train: 训练集特征
    :param y_train: 训练集目标值
    :param max_depth: 树的最大深度，默认值为3
    :param n_estimators: 弱学习器数量，默认值为100
    :param show_progress: 是否显示训练进度和摘要，默认True
    :return: 训练好的模型实例
    """
    if show_progress:
        print(f"训练数据量: {X_train.shape[0]} 个样本, {X_train.shape[1]} 个特征")

    # 进一步划分训练集为训练和验证集
    X_train_final, X_val, y_train_final, y_val = train_test_split(
        X_train, y_train, test_size=0.2, random_state=42
    )

    # 创建XGBoost回归器实例
    reg = xgb.XGBRegressor(
        objective="reg:squarederror",
        max_depth=max_depth,
        n_estimators=n_estimators,
        eval_metric=[
            "rmse",  # 评估指标：均方根误差（RMSE）
        ],
        early_stopping_rounds=10,  # 验证集上RMSE未改善10轮时停止训练
        verbosity=1,
    )

    # 使用训练数据拟合模型
    reg.fit(
        X_train_final,
        y_train_final,
        eval_set=[(X_train_final, y_train_final), (X_val, y_val)],
    )

    if show_progress:
        # 显示训练摘要
        print("\n=== 训练摘要 ===")
        best_iteration = (
            reg.best_iteration if hasattr(reg, "best_iteration") else n_estimators
        )
        print(f"最佳迭代次数: {best_iteration}")
        print(f"最终训练RMSE: {reg.evals_result_['validation_0']['rmse'][-1]:.4f}")
        print(f"最终验证RMSE: {reg.evals_result_['validation_1']['rmse'][-1]:.4f}")

        # 绘制训练历史
        plot_training_history(reg.evals_result_)

    return reg


def evaluate_model(model, X_test, y_test):
    """
    评估模型性能
    :param model: 训练好的模型实例
    :param X_test: 测试集特征
    :param y_test: 测试集目标值
    :return: 均方根误差（RMSE）
    """
    # 使用测试集进行预测，并计算均方根误差（RMSE）
    rmse = root_mean_squared_error(y_test, model.predict(X_test))
    # RMSE是回归问题常用的评估指标，值越小表示模型预测越准确
    print("RMSE:", rmse)
    return rmse


def save_model(model, model_path):
    """
    保存模型到指定路径
    :param model: 训练好的模型实例
    :param model_path: 模型保存路径
    """
    if not model_path.endswith(".model"):
        model_path += ".model"
    model.save_model(model_path)
    print(f"模型已保存至 {model_path}")


def plot_training_history(history):
    """
    绘制训练历史曲线
    :param history: 训练历史数据
    """
    plt.figure(figsize=(12, 4))

    # 训练集RMSE
    plt.subplot(1, 2, 1)
    plt.plot(history["validation_0"]["rmse"], label="Train RMSE")
    plt.plot(history["validation_1"]["rmse"], label="Validation RMSE")
    plt.xlabel("Epoch")
    plt.ylabel("RMSE")
    plt.title("RMSE during Training")
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.show()


# 主程序执行
if __name__ == "__main__":
    # 训练模型
    print("开始训练模型...")
    model = train_model(X_train, y_train)

    # 评估模型
    print("模型评估结果:")
    test_rmse = evaluate_model(model, X_test, y_test)

    # # 保存模型
    # save_model(model, "models/diabetes_xgboost_model.model")
